'''
配置文件
'''

import os
import torch
import datetime
import logging
import flair
from easydict import EasyDict as edict  # easydict的作用：可以使得以属性的方式去访问字典的值


# 初始化两个字典
__C = edict()
args = __C

# Base
__C.gpu_id = 5  #建议使用环境变量控制要使用的GPU   export CUDA_VISIBLE_DEVICES=6
__C.num_workers = 1
__C.seed = 666
__C.do_train = True
__C.do_test = False
__C.batch_size = 16
__C.lr = 0.0001
__C.now = datetime.datetime.now()
__C.device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.cuda.set_device(__C.gpu_id)   # 建议使用环境变量控制要使用的GPU   export CUDA_VISIBLE_DEVICES=6
flair.device = torch.device(__C.gpu_id)

# dataset options
__C.dataset = edict()
__C.dataset.name = 'twitter2017'
__C.dataset.type = 'none'
__C.dataset.data_dir = os.path.join('dataset',  __C.dataset.name)       # dateset/twitter2017
__C.dataset.save_dir = os.path.join('result', __C.dataset.name, 'predict')
__C.dataset.ckpt_dir = os.path.join('result', __C.dataset.name, 'ckpt')
__C.dataset.log_dir = os.path.join('result', __C.dataset.name, 'log')

__C.dataset.predict_file = __C.dataset.save_dir + "/{}/epoch{}.txt"
__C.dataset.image_caption = os.path.join('captions', '{}_image_captions.json')
__C.dataset.text = os.path.join('texts', '{}_text.json')
__C.dataset.image = os.path.join('images', 'image_features_{}.pkl')
__C.dataset = dict(__C.dataset)


# 创建文件夹
if not os.path.exists(__C.dataset.save_dir):
    os.makedirs(__C.dataset.save_dir)
if not os.path.exists(__C.dataset.ckpt_dir):
    os.makedirs(__C.dataset.ckpt_dir)
if not os.path.exists(__C.dataset.log_dir):
    os.makedirs(__C.dataset.log_dir)

# logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)-8s %(message)s')
logFormatter = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s')
rootLogger = logging.getLogger()
fileHandler = logging.FileHandler(
        os.path.join(__C.dataset.log_dir, 'stdout_{}.{}_{}.{}.log'.format(__C.now.month, __C.now.day, __C.now.hour, __C.now.minute)), 'w+')
fileHandler.setFormatter(logFormatter)
rootLogger.addHandler(fileHandler)
__C.logging = logging

# training options
__C.train = edict()
__C.train.max_epochs = 100
__C.train.max_word_length = 128
__C.train.bert_dim = 4246
__C.train.module_dim = 2048
__C.train.interval = 20
__C.train.tag2idx = {
    "PAD": 0,
    "B-PER": 1,
    "I-PER": 2,
    "B-LOC": 3,
    "I-LOC": 4,
    "B-ORG": 5,
    "I-ORG": 6,
    "B-MISC": 7,
    "I-MISC": 8,
    "O": 9,
    "X": 10
}
__C.train.idx2tag = {str(idx): tag for tag, idx in __C.train.tag2idx.items()}
__C.train.pic_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
           'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
           'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
           'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
           'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
           'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
           'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
           'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
           'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
           'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
           'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
__C.train = dict(__C.train)
